"""
NVIDIA显卡静态设备信息采集脚本
采集显卡型号、显存容量、算力等不变的硬件规格信息
"""

import subprocess
import json
from typing import List, Optional


class NVIDIAGPUStaticCollector:
    def __init__(self):
        # GPU规格数据库 - 存储各型号显卡的固定硬件参数
        self.gpu_specs = {
            'RTX 4090': {
                'cuda_cores': 16384,  # CUDA核心数
                'base_clock': 2235,  # 基础频率 MHz
                'boost_clock': 2520,  # 加速频率 MHz
                'tflops_fp32': 83.0,  # 单精度浮点算力 TFLOPS
                'tflops_fp16': 166.0  # 半精度浮点算力 TFLOPS
            },
            'H100': {
                'cuda_cores': 14592,
                'base_clock': 1095,
                'boost_clock': 1980,
                'tflops_fp32': 60.0,
                'tflops_fp16': 120.0
            }
        }

    def parse_gpu_name(self, raw_name: str) -> str:
        """解析GPU型号名称"""
        name = raw_name.replace('NVIDIA ', '').replace('GeForce ', '').strip()

        if '4090' in name:
            return 'RTX 4090'
        elif 'H100' in name:
            return 'H100'

        return name

    def get_gpu_static_info(self, gpu_index: int) -> Optional[dict]:
        """获取单个GPU的静态硬件信息"""
        try:
            # 查询GPU基本硬件信息
            cmd = [
                'nvidia-smi', '-i', str(gpu_index),
                '--query-gpu=name,memory.total,pcie.link.gen.max,pcie.link.width.max',
                '--format=csv,noheader,nounits'
            ]

            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            data = result.stdout.strip().split(', ')

            if len(data) < 4:
                return None

            # 解析GPU型号并获取规格
            gpu_name = self.parse_gpu_name(data[0])
            specs = self.gpu_specs.get(gpu_name, {})

            # 安全转换数值
            def safe_int(value):
                try:
                    return int(value) if value != '[Not Supported]' else 0
                except:
                    return 0

            memory_total = safe_int(data[1])

            gpu_info = {
                'index': gpu_index,  # GPU索引
                'name': gpu_name,  # GPU名称
                'specifications': {  # 规格参数
                    'cuda_cores': specs.get('cuda_cores', 0),  # CUDA核心数
                    'base_clock_mhz': specs.get('base_clock', 0),  # 基础频率
                    'boost_clock_mhz': specs.get('boost_clock', 0),  # 加速频率
                    'memory_total_mb': memory_total,  # 总显存 MB
                    'memory_total_gb': round(memory_total / 1024, 1),  # 总显存 GB
                    'pcie_gen': safe_int(data[2]),  # PCIe代数
                    'pcie_width': safe_int(data[3])  # PCIe通道宽度
                },
                'compute_capability': {  # 计算能力
                    'tflops_fp32': specs.get('tflops_fp32', 0.0),  # 单精度算力
                    'tflops_fp16': specs.get('tflops_fp16', 0.0)  # 半精度算力
                }
            }

            return gpu_info

        except subprocess.CalledProcessError:
            return None

    def calculate_static_summary(self, gpus: List[dict]) -> dict:
        """计算静态信息汇总统计"""
        if not gpus:
            return {
                'gpu_count': 0,
                'gpu_models': {},
                'total_specs': {
                    'total_cuda_cores': 0,
                    'total_tflops_fp32': 0.0,
                    'total_tflops_fp16': 0.0,
                    'total_memory_mb': 0,
                    'total_memory_gb': 0.0
                }
            }

        # 统计各型号GPU数量
        gpu_models = {}
        for gpu in gpus:
            model = gpu['name']
            gpu_models[model] = gpu_models.get(model, 0) + 1

        # 计算总规格
        total_cuda_cores = sum(gpu['specifications']['cuda_cores'] for gpu in gpus)
        total_tflops_fp32 = sum(gpu['compute_capability']['tflops_fp32'] for gpu in gpus)
        total_tflops_fp16 = sum(gpu['compute_capability']['tflops_fp16'] for gpu in gpus)
        total_memory_mb = sum(gpu['specifications']['memory_total_mb'] for gpu in gpus)

        return {
            'gpu_count': len(gpus),  # GPU总数
            'gpu_models': gpu_models,  # 各型号数量统计
            'total_specs': {  # 总体规格
                'total_cuda_cores': total_cuda_cores,  # 总CUDA核心数
                'total_tflops_fp32': round(total_tflops_fp32, 1),  # 总单精度算力
                'total_tflops_fp16': round(total_tflops_fp16, 1),  # 总半精度算力
                'total_memory_mb': total_memory_mb,  # 总显存MB
                'total_memory_gb': round(total_memory_mb / 1024, 1)  # 总显存GB
            }
        }

    def collect_static_info(self) -> dict:
        """收集所有GPU静态信息"""
        try:
            # 检查nvidia-smi是否可用
            subprocess.run(['nvidia-smi', '-L'], capture_output=True, check=True)

            # 获取GPU列表
            result = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True, check=True)
            gpu_lines = result.stdout.strip().split('\n')

            if not gpu_lines or gpu_lines == ['']:
                return {
                    'status': 'no_gpu',  # 状态：无GPU
                    'message': '未检测到NVIDIA显卡',
                    'summary': self.calculate_static_summary([]),
                    'gpus': []
                }

            # 获取所有GPU静态信息
            gpus = []
            for i in range(len(gpu_lines)):
                gpu_info = self.get_gpu_static_info(i)
                if gpu_info:
                    gpus.append(gpu_info)

            # 计算汇总统计
            summary = self.calculate_static_summary(gpus)

            return {
                'status': 'success',  # 状态：成功
                'message': f'成功检测到 {len(gpus)} 张显卡',
                'summary': summary,  # 汇总信息
                'gpus': gpus  # 详细GPU列表
            }

        except (subprocess.CalledProcessError, FileNotFoundError):
            return {
                'status': 'no_driver',  # 状态：无驱动
                'message': 'NVIDIA驱动未安装或nvidia-smi不可用',
                'summary': self.calculate_static_summary([]),
                'gpus': []
            }


def main():
    collector = NVIDIAGPUStaticCollector()
    result = collector.collect_static_info()
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()
